# 导入库
import random
import numpy as np
import csv
import time
from tqdm import tqdm
from RA_Cal.Costfunction import Costfunction_tt


class DPGA:
    def __init__(self, fitness_func, n, population_size, generations, casei,
                 div_population_ratio=0.5, migration_rate=0.1, migration_interval=5,
                 diversity_threshold=0.3, elite_size=2):
        # 外部传参
        self.fitness_func = fitness_func  # 适应度评估函数
        self.n = n  # 编码长度
        self.population_size = population_size  # 主群体中个体数量
        self.generations = generations  # 迭代次数
        self.casei = casei

        # 双种群参数
        self.div_population_size = int(100)  # 多样性种群大小
        # self.div_population_size = int(population_size * div_population_ratio)  # 多样性种群大小
        self.migration_rate = migration_rate  # 迁移率
        self.migration_interval = migration_interval  # 迁移间隔
        self.diversity_threshold = diversity_threshold  # 多样性阈值
        self.elite_size = elite_size  # 精英个体数量

        # 内部传参
        self.main_population = np.zeros((self.population_size, self.n))  # 主种群
        self.div_population = np.zeros((self.div_population_size, self.n))  # 多样性种群
        self.y_main_population = np.zeros(self.population_size)  # 主种群适应度
        self.y_div_population = np.zeros(self.div_population_size)  # 多样性种群适应度

        # 自适应参数
        self.base_mutation_rate = 0.01  # 基础变异率
        self.base_crossover_rate = 0.8  # 基础交叉率
        self.current_main_diversity = 0  # 当前主种群多样性
        self.current_div_diversity = 0  # 当前多样性种群多样性

        # 精英保留
        self.main_elite = None  # 主种群精英
        self.div_elite = None  # 多样性种群精英

        self.initialize_populations()  # 初始化双种群
        self.best_fitness = -np.inf  # 初始化最佳适应度为无穷小量
        self.best_solution = None  # 最佳解
        self.number = 0  # 初始化迭代次数

    # 初始化双种群
    def initialize_populations(self):
        self.main_population = np.random.randint(0, 2, (self.population_size, self.n))
        self.div_population = np.random.randint(0, 2, (self.div_population_size, self.n))

    # 计算汉明距离
    def hamming_distance(self, ind1, ind2):
        return np.sum(ind1 != ind2)

    # 计算种群多样性
    def calculate_diversity(self, population):
        n = len(population)
        if n <= 1:
            return 0
        # 使用广播计算所有个体间的汉明距离
        distances = np.sum(population[:, np.newaxis, :] != population[np.newaxis, :, :], axis=2)
        # 只取上三角矩阵的平均值
        return np.sum(np.triu(distances, k=1)) / (n * (n - 1) / 2)

    # 计算自适应变异率
    def adaptive_mutation_rate(self, population_type='main'):
        diversity = self.current_main_diversity if population_type == 'main' else self.current_div_diversity

        # 当多样性低时增加变异率，反之降低
        if diversity < self.diversity_threshold:
            return min(0.1, self.base_mutation_rate * (1 + (self.diversity_threshold - diversity)))
        else:
            return max(0.001, self.base_mutation_rate * (1 - (diversity - self.diversity_threshold)))

    # 计算自适应交叉率
    def adaptive_crossover_rate(self, population_type='main'):
        diversity = self.current_main_diversity if population_type == 'main' else self.current_div_diversity

        # 当多样性低时降低交叉率，反之增加
        if diversity < self.diversity_threshold:
            return max(0.6, self.base_crossover_rate * (1 - (self.diversity_threshold - diversity)))
        else:
            return min(0.95, self.base_crossover_rate * (1 + (diversity - self.diversity_threshold)))

    # 保存精英个体
    def save_elites(self):
        # 主种群精英
        main_elite_indices = np.argsort(self.y_main_population)[-self.elite_size:]
        main_elite_indices = main_elite_indices.astype(int)  # 确保是整数数组
        self.main_elite = self.main_population[main_elite_indices].copy()

        # 多样性种群精英（结合适应度和多样性）
        div_scores = np.zeros(self.div_population_size)
        for i in range(self.div_population_size):
            # 计算多样性得分
            diversity_score = np.mean([self.hamming_distance(self.div_population[i], other)
                                       for other in self.div_population])
            fitness_score = self.y_div_population[i]

            # 安全的归一化计算
            max_diversity = max(1e-10, np.max(diversity_score))  # 避免除零
            max_fitness = max(1e-10, np.max(self.y_div_population))  # 避免除零

            # 如果适应度全为负，则将其转换为非负值
            if max_fitness < 0:
                normalized_fitness = (fitness_score - np.min(self.y_div_population)) / \
                                     (np.max(self.y_div_population) - np.min(self.y_div_population) + 1e-10)
            else:
                normalized_fitness = fitness_score / max_fitness

            # 计算综合得分
            div_scores[i] = 0.5 * (diversity_score / max_diversity) + 0.5 * normalized_fitness

        div_elite_indices = np.argsort(div_scores)[-self.elite_size:]
        div_elite_indices = div_elite_indices.astype(int)  # 确保是整数数组
        self.div_elite = self.div_population[div_elite_indices].copy()

    # 恢复精英个体
    def restore_elites(self):
        if self.main_elite is not None:
            # 替换主种群中最差的个体
            worst_indices = np.argsort(self.y_main_population)[:self.elite_size]
            self.main_population[worst_indices] = self.main_elite

        if self.div_elite is not None:
            # 替换多样性种群中最差的个体
            worst_indices = np.argsort(self.y_div_population)[:self.elite_size]
            self.div_population[worst_indices] = self.div_elite

    # 改进的迁移策略
    def migrate(self):
        if self.number % self.migration_interval == 0:
            num_migrants = int(self.population_size * self.migration_rate)

            # 动态调整迁移数量基于多样性
            if self.current_main_diversity < self.diversity_threshold:
                num_migrants = int(num_migrants * 1.5)  # 增加迁移量
            elif self.current_div_diversity < self.diversity_threshold:
                num_migrants = int(num_migrants * 0.5)  # 减少迁移量

            # 从主种群选择最优个体迁移到多样性种群
            main_best_indices = np.argsort(self.y_main_population)[-num_migrants:]
            div_worst_indices = np.argsort(self.y_div_population)[:num_migrants]

            for i in range(min(len(main_best_indices), len(div_worst_indices))):
                self.div_population[div_worst_indices[i]] = self.main_population[main_best_indices[i]].copy()

            # 从多样性种群选择最具多样性的个体迁移到主种群
            div_diversity = np.array([
                np.mean([self.hamming_distance(ind, other)
                         for other in self.div_population])
                for ind in self.div_population
            ])
            div_best_indices = np.argsort(div_diversity)[-num_migrants:]
            main_worst_indices = np.argsort(self.y_main_population)[:num_migrants]

            for i in range(min(len(div_best_indices), len(main_worst_indices))):
                self.main_population[main_worst_indices[i]] = self.div_population[div_best_indices[i]].copy()

    # 改进的变异操作
    def mutation(self):
        # 主种群变异
        main_mutation_rate = self.adaptive_mutation_rate('main')
        mutation_mask_main = np.random.random((self.population_size, self.n)) < main_mutation_rate
        self.main_population = np.logical_xor(self.main_population, mutation_mask_main)

        # 多样性种群变异
        div_mutation_rate = self.adaptive_mutation_rate('div')
        mutation_mask_div = np.random.random((self.div_population_size, self.n)) < div_mutation_rate
        self.div_population = np.logical_xor(self.div_population, mutation_mask_div)

    # 改进的交叉操作
    def crossover(self):
        # 主种群交叉
        main_crossover_rate = self.adaptive_crossover_rate('main')
        crossover_mask = np.random.random(self.population_size // 2) < main_crossover_rate

        # 随机配对
        indices = np.arange(self.population_size)
        np.random.shuffle(indices)
        pairs = indices.reshape(-1, 2)

        # 生成交叉点
        points = np.sort(np.random.randint(1, self.n, (len(pairs), 2)), axis=1)

        # 执行交叉
        for i, (p1, p2) in enumerate(pairs):
            if crossover_mask[i]:
                start, end = points[i]
                temp = self.main_population[p1, start:end].copy()
                self.main_population[p1, start:end] = self.main_population[p2, start:end]
                self.main_population[p2, start:end] = temp

    # 计算适应度并保存
    def cal_fitness(self):
        # 向量化计算适应度
        self.y_main_population = np.array([self.fitness_func(self.casei, x) for x in self.main_population])
        self.y_div_population = np.array([self.fitness_func(self.casei, x) for x in self.div_population])

        # 更新全局最优
        current_best = max(np.max(self.y_main_population), np.max(self.y_div_population))
        if current_best > self.best_fitness:
            self.best_fitness = current_best
            with open("DPGA_result.csv", 'a', encoding='utf-8', newline='') as f:
                write = csv.writer(f)
                write.writerow([self.number, self.best_fitness])

    # 选择操作
    def selection(self):
        # 主种群选择
        competitors_main = np.random.randint(0, self.population_size, (self.population_size, 3))
        competitors_fitness_main = self.y_main_population[competitors_main]
        winners_main = competitors_main[np.arange(len(competitors_main)),
        np.argmax(competitors_fitness_main, axis=1)]
        self.main_population = self.main_population[winners_main]

        # 多样性种群选择
        competitors_div = np.random.randint(0, self.div_population_size, (self.div_population_size, 3))
        # 计算竞争者的多样性得分
        diversity_matrix = np.sum(self.div_population[:, np.newaxis, :] !=
                                  self.div_population[np.newaxis, :, :], axis=2)
        diversity_scores = np.mean(diversity_matrix, axis=1)

        # 归一化处理
        norm_diversity = diversity_scores / (np.max(diversity_scores) + 1e-10)
        norm_fitness = self.y_div_population / (np.max(self.y_div_population) + 1e-10)

        # 计算综合得分
        combined_scores = 0.5 * norm_diversity[competitors_div] + 0.5 * norm_fitness[competitors_div]
        winners_div = competitors_div[np.arange(len(competitors_div)),
        np.argmax(combined_scores, axis=1)]
        self.div_population = self.div_population[winners_div]

    # 改进的运行方法
    def run(self):
        for _ in range(self.generations):
            sti = time.time()
            self.number += 1

            # 保存当前精英
            self.save_elites()

            # 评估适应度
            self.cal_fitness()

            # 计算当前多样性
            self.current_main_diversity = self.calculate_diversity(self.main_population)
            self.current_div_diversity = self.calculate_diversity(self.div_population)

            # 种群迁移
            self.migrate()

            # 选择、交叉、变异操作
            self.selection()
            self.crossover()
            self.mutation()

            # 恢复精英
            self.restore_elites()

            # 输出详细信息
            print(f"Iter {_}: Best Fitness={self.best_fitness:.4f}, "
                  f"Main Diversity={self.current_main_diversity:.4f}, "
                  f"Div Diversity={self.current_div_diversity:.4f}, "
                  f"Main Mutation Rate={self.adaptive_mutation_rate('main'):.4f}, "
                  f"Time={time.time() - sti:.2f}s")

        # 获取最终最优解
        all_solutions = np.vstack([self.main_population, self.div_population])
        all_fitness = np.concatenate([self.y_main_population, self.y_div_population])
        best_idx = np.argmax(all_fitness)
        best_code = all_solutions[best_idx]

        # 保存结果
        np.set_printoptions(threshold=np.inf)
        with open("DPGA_result.csv", 'a', encoding='utf-8', newline='') as f:
            write = csv.writer(f)
            write.writerow([self.best_fitness, best_code])

        return best_code, self.best_fitness
